/* * Copyright © 2023 https://www.cestc.cn/ All rights reserved. */
package com.ffz.calcitesql.sqlparser.vistor;

import com.ffz.calcitesql.sqlparser.enums.LineNodeType;
import com.ffz.calcitesql.sqlparser.enums.SqlOperatorType;
import com.ffz.calcitesql.sqlparser.meta.Column;
import com.ffz.calcitesql.sqlparser.meta.LineColumnNode;
import com.ffz.calcitesql.sqlparser.meta.Table;
import com.google.common.collect.ImmutableList;
import org.apache.calcite.sql.*;
import org.apache.calcite.sql.util.SqlBasicVisitor;
import org.apache.commons.lang.StringUtils;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * 访问者模式解析.
 *
 * @author fengfazheng
 * @since 2024/1/16 17:43
 */
public class SqlFieldVisitor extends SqlBasicVisitor<Void> {
    private final static String INIT_TABLE = "RS";
    public LinkedList<LineColumnNode> lineColumnNodeList;
    public Map<String, LinkedList<LineColumnNode>> lineColumnNodeMap = new HashMap<>();

    @Override
    public Void visit(SqlIdentifier id) {
        return super.visit(id);
    }

    @Override
    public Void visit(SqlCall sqlCall) {
        if (sqlCall instanceof SqlSelect) {
            // 处理 SELECT 语句
            SqlSelect select = (SqlSelect) sqlCall;
            hanleSqlSelectList(select.getSelectList());
            // 血缘关系继续往上查找
            handleSqlSelectFrom(select.getFrom());
        } else if (sqlCall instanceof SqlBasicCall) {
            handlerSqlBasicCall((SqlBasicCall) sqlCall);
        }
        return super.visit(sqlCall);
    }

    private void handlerSqlBasicCall(SqlBasicCall sqlBasicCall) {
        SqlKind kind = sqlBasicCall.getKind();
        switch (kind) {
            case AS:
                List<SqlNode> operandList = sqlBasicCall.getOperandList();
                System.out.println("===table name===" + operandList.get(0) + "===alias===" + operandList.get(1));
                break;
            default:
                break;
        }
    }

    private void hanleSqlSelectList(SqlNodeList sqlNodeList) {
        for (SqlNode sqlNode : sqlNodeList) {
            lineColumnNodeList = new LinkedList<>();
            // 遍历查询节点，构造血缘
            if (sqlNode instanceof SqlIdentifier) {
                handleSelectSqlIdentifier((SqlIdentifier) sqlNode);
            } else if (sqlNode instanceof SqlBasicCall) {
                SqlBasicCall sqlBasicCall = (SqlBasicCall) sqlNode;
                SqlOperator operator1 = sqlBasicCall.getOperator();
                List<SqlNode> operandList = sqlBasicCall.getOperandList();
                // 处理带别名情况（X.Y AS Z）(SUM(X.Y) AS Z)
                if (operator1 instanceof SqlAsOperator) {
                    LineColumnNode first = LineColumnNode.builder()
                            .name(operandList.get(1).toString()).type(LineNodeType.COLUMN.name()).table(INIT_TABLE).build();
                    lineColumnNodeList.addFirst(first);
                    SqlNode orginSqlNode = operandList.get(0);
                    if (orginSqlNode instanceof SqlIdentifier) {
                        // 正常字段别名情况
                        System.out.println(orginSqlNode.toString());
                    } else if (orginSqlNode instanceof SqlBasicCall) {
                        // 为函数等情况
                        SqlBasicCall sqlBasicCall1 = (SqlBasicCall) orginSqlNode;
                        SqlOperator operator = sqlBasicCall1.getOperator();
                        if (operator instanceof SqlUnresolvedFunction) {
                            handleSelectFunction(sqlBasicCall1);
                        }
                    }
                    lineColumnNodeMap.put(operandList.get(1).toString(), lineColumnNodeList);
                } else if (operator1 instanceof SqlUnresolvedFunction) {
                    // 处理不带别名函数（SUM(X.Y)）
                    SqlNode sqlNode1 = sqlBasicCall.getOperandList().get(0);
                    LineColumnNode first =
                            LineColumnNode.builder().type(LineNodeType.FUNCTION.name())
                                    .funciton(operator1.getName()).name(sqlNode1.toString()).build();
                    lineColumnNodeList.addFirst(first);
                    if (sqlNode1 instanceof SqlBasicCall) {
                        // 处理函数嵌套问题
                    }
                    lineColumnNodeMap.put(sqlNode1.toString(), lineColumnNodeList);
                }

            } else if (sqlNode instanceof SqlNumericLiteral) {
                SqlNumericLiteral sqlNumericLiteral = (SqlNumericLiteral) sqlNode;
                String num = sqlNumericLiteral.toValue();
                LineColumnNode first =
                        LineColumnNode.builder().type(LineNodeType.COLUMN.name()).name(num).table(INIT_TABLE).build();
                lineColumnNodeList.addFirst(first);
                lineColumnNodeList.addFirst(buildHeadNode());
                lineColumnNodeMap.put(num, lineColumnNodeList);
            }
        }
    }

    private void handleSelectSqlIdentifier(SqlIdentifier selectFeild) {
        if ("*".equalsIgnoreCase(selectFeild.toString())) {
            LineColumnNode first =
                    LineColumnNode.builder().name(selectFeild.toString())
                            .type(LineNodeType.ALL_COLUMN.name()).table(INIT_TABLE).build();
            lineColumnNodeList.addFirst(first);
            lineColumnNodeMap.put(selectFeild.toString(), lineColumnNodeList);
        } else {
            if(selectFeild.names.size()>=2){
                // 带表名
                String[] split = selectFeild.toString().split("\\.");
                // 构造初始节点
                LineColumnNode first =
                        LineColumnNode.builder().name(split[1]).type(LineNodeType.COLUMN.name()).table(INIT_TABLE).build();
                LineColumnNode next = LineColumnNode.builder().name(split[1])
                        .type(LineNodeType.COLUMN.name()).table(split[0]).build();
                lineColumnNodeList.addFirst(first);
                lineColumnNodeList.addFirst(next);
                lineColumnNodeMap.put(split[1], lineColumnNodeList);
            }else {
                // 不带表名
                LineColumnNode first =
                        LineColumnNode.builder().name(selectFeild.toString()).type(LineNodeType.COLUMN.name()).table(INIT_TABLE).build();
                lineColumnNodeList.addFirst(first);
                lineColumnNodeMap.put(selectFeild.toString(), lineColumnNodeList);
            }

        }
    }

    private void handleSelectFunction(SqlBasicCall sqlBasicCall) {
        // 设置函数节点
        SqlNode sqlNode1 = sqlBasicCall.getOperandList().get(0);
        LineColumnNode second =
                LineColumnNode.builder().type(LineNodeType.FUNCTION.name())
                        .funciton(sqlBasicCall.getOperator().getName()).build();
        lineColumnNodeList.addFirst(second);
        if (sqlNode1 instanceof SqlIdentifier) {
            SqlIdentifier sqlIdentifier = (SqlIdentifier) sqlNode1;
            if (sqlIdentifier.names.size() >= 2) {
                LineColumnNode third =
                        LineColumnNode.builder().type(LineNodeType.COLUMN.name())
                                .name(sqlIdentifier.names.get(1)).table(sqlIdentifier.names.get(0)).build();
                lineColumnNodeList.addFirst(third);
            } else {
                LineColumnNode third =
                        LineColumnNode.builder().type(LineNodeType.COLUMN.name())
                                .name(sqlIdentifier.names.get(0)).build();
                lineColumnNodeList.addFirst(third);
            }
        } else if (sqlNode1 instanceof SqlBasicCall) {
            // TODO 处理函数套函数情况
        }
    }

    private void handleSqlSelectFrom(SqlNode sqlNode) {
        if (sqlNode instanceof SqlJoin) {
            handlefromJoin((SqlJoin) sqlNode);
        } else if (sqlNode instanceof SqlBasicCall) {
            Map<String, Table> tableAsNameMap = new HashMap<>();
            // 单表场景
            SqlBasicCall sqlBasicCall = (SqlBasicCall) sqlNode;
            SqlOperator operator = sqlBasicCall.getOperator();
            if (operator instanceof SqlAsOperator) {
                List<SqlNode> operandList = sqlBasicCall.getOperandList();
                SqlNode tableSqlNode = operandList.get(0);
                Table table = null;
                SqlNode asSqlNode = operandList.get(1);
                if (tableSqlNode instanceof SqlIdentifier) {
                    // 获取表名/库名
                    table = buildAsTable((SqlIdentifier) tableSqlNode, asSqlNode.toString());
                } else {
                    // TODO 子查询等其他情况
                }
                tableAsNameMap.put(asSqlNode.toString(), table);
            }
            // 遍历已有血缘并添加节点
            visitAddLineColumnNode(tableAsNameMap);
        } else if (sqlNode instanceof SqlIdentifier) {
            handleTableSqlIdentifier((SqlIdentifier) sqlNode);
        }
    }

    private Table buildAsTable(SqlIdentifier sqlIdentifier, String alias) {
        // 设置表名/库名
        String dbName = null;
        String tableName;
        ImmutableList<String> names = sqlIdentifier.names;
        if (names.size() >= 2) {
            dbName = names.get(0);
            tableName = names.get(1);
        } else {
            tableName = names.get(0);
        }
        return Table.builder().name(tableName)
                .db(dbName)
                .alias(alias).build();
    }

    private void handleTableSqlIdentifier(SqlIdentifier sqlIdentifier) {
        Table table = buildAsTable(sqlIdentifier, null);
        vistSingleTableWithNoAsOperator(table);
    }

    private void vistSingleTableWithNoAsOperator(Table table) {
        for (Map.Entry<String, LinkedList<LineColumnNode>> entry : lineColumnNodeMap.entrySet()) {
            LinkedList<LineColumnNode> value = entry.getValue();
            // 取头节点
            LineColumnNode lineColumnNode = value.get(0);
            if (LineNodeType.HEAD.name().equalsIgnoreCase(lineColumnNode.toString())) {
                // 头部节点过滤
                continue;
            }
            LineColumnNode name =
                    LineColumnNode.builder().type(LineNodeType.COLUMN.name())
                            .table(table.getName()).name(lineColumnNode.getName()).build();
            value.addFirst(name);
            // 字段无法再向上追溯,添加头部节点
            value.addFirst(buildHeadNode());
        }
    }

private void visitAddLineColumnNode(Map<String, Table> tableAsNameMap) {
    for (Map.Entry<String, LinkedList<LineColumnNode>> entry : lineColumnNodeMap.entrySet()) {
        String key = entry.getKey();
        LinkedList<LineColumnNode> value = entry.getValue();
        // 取头节点
        LineColumnNode lineColumnNode = value.get(0);
        if (LineNodeType.HEAD.name().equalsIgnoreCase(lineColumnNode.toString())) {
            // 头部节点过滤
            continue;
        }
        String table = lineColumnNode.getTable();
        // 表名不为空且集合内存在，即添加头节点
        if (StringUtils.isNotEmpty(table)) {
            Table table1 = tableAsNameMap.get(table);
            if (table1 != null) {
                // 字段向上追溯表名
                LineColumnNode name =
                        LineColumnNode.builder().type(LineNodeType.COLUMN.name())
                                .table(table1.getName()).name(lineColumnNode.getName()).build();
                value.addFirst(name);
            } else {
                // 字段无法再向上追溯,添加头部节点
                value.addFirst(buildHeadNode());
            }

        }
    }
}

private LineColumnNode buildHeadNode() {
    return LineColumnNode.builder().type(LineNodeType.HEAD.name()).build();
}

private void handlefromJoin(SqlJoin sqlJoin) {
    SqlNode left = sqlJoin.getLeft();
    Map<String, Table> tableAsNameMap = new HashMap<>();
    if (left instanceof SqlBasicCall) {
        SqlBasicCall leftSqlBasicCall = (SqlBasicCall) left;
        if (leftSqlBasicCall.getKind().equals(SqlKind.AS)) {
            List<SqlNode> operandList = leftSqlBasicCall.getOperandList();
            SqlNode sqlNode = operandList.get(0);
            SqlNode asSqlNode = operandList.get(1);
            Table table = Table.builder().name(sqlNode.toString()).alias(asSqlNode.toString()).build();
            tableAsNameMap.put(asSqlNode.toString(), table);
        }
    }
    SqlNode right = sqlJoin.getRight();
    if (right instanceof SqlBasicCall) {
        SqlBasicCall rightSqlBasicCall = (SqlBasicCall) right;
        if (rightSqlBasicCall.getKind().equals(SqlKind.AS)) {
            List<SqlNode> operandList = rightSqlBasicCall.getOperandList();
            SqlNode sqlNode = operandList.get(0);
            SqlNode asSqlNode = operandList.get(1);
            Table table = Table.builder().name(sqlNode.toString()).alias(asSqlNode.toString()).build();
            tableAsNameMap.put(asSqlNode.toString(), table);
        }
    }
    // 遍历已有血缘并添加节点
    visitAddLineColumnNode(tableAsNameMap);
}
}
